import pathlib

import numba
import numpy

from problem2.checkpoint import Checkpoint
from utilities.checkpoint_manager import CheckpointManager


@numba.jit(nopython=True)
def _scan_numba(calculation_state, represented_points):
    points = numpy.unique(represented_points)
    result = numpy.empty((len(points), 2))
    for i in range(len(points)):
        indexes = represented_points == points[i]
        result[i, 0] = points[i] * 0.02
        states = calculation_state[indexes] * 0.02
        states = numpy.sin(numpy.pi * states)
        result[i, 1] = numpy.average(states)
    return result


def main(checkpoint_manager: CheckpointManager):
    for i, checkpoint_path in checkpoint_manager.existed().items():
        print(i)
        checkpoint = Checkpoint.load(checkpoint_path)
        result = _scan_numba(checkpoint.calculation_state(), checkpoint.represented_points())
        numpy.save(checkpoint_path.with_suffix(".npy"), result)


if __name__ == "__main__":
    main(CheckpointManager(pathlib.Path("./outputs")))
